#include <iostream>
#include <iomanip>
#include <vector>
#include <cstdlib>
#include <omp.h>
#include <complex>
#include <mambo.hpp>
#include <dns_ops.h>
using namespace std;

int main(int argc, char** argv)
{

  typedef  double real_type;
  typedef  complex<double> complex_type;
  const int n_var=3;
  int niters=2;
  int nz=12288;
  int npx=512;
  int npy=512;
  bool debug=false;
  int ic=1;
  while(ic<argc)
  {
    string a(argv[ic]);
    if(a.find("-i")<a.size()) 
      niters=atoi(argv[++ic]);
    else if(a.find("-s")<a.size()) //get the z
      nz=atoi(argv[++ic]);
    else if(a.find("-npx")<a.size()) //number of pencils
      npx=atoi(argv[++ic]);
    else if(a.find("-npy")<a.size()) //number of pencils
      npy=atoi(argv[++ic]);
    else if(a.find("-debug")<a.size())  
      debug=true;
    else if(a.find("-h")<a.size())
    {
      cerr << "Usage:\n  rk4 -s fft-size -npx x-grid -npy y-grid -i iterations" << endl;
      cerr << "Defaults: \n rk4 -s 12288 -npx 512 -npy 512 -i 2 " << endl;
    }
    ++ic;
  }


  cout << "BW nx_loc=" << nz/npx << " ny_loc=" << nz/npy << " nz_loc=" << nz
    << endl;

  int nc=nz/npx*nz/npy;
  int ntot=n_var*nz*nc;
  int np=omp_get_max_threads();

#if defined(USE_MALLOC)
  real_type *Q = static_cast<real_type*>(malloc(np*ntot*sizeof(real_type)));
  real_type *rhs = static_cast<real_type*>(malloc(np*ntot*sizeof(real_type)));
  real_type *Q_tmp = static_cast<real_type*>(malloc(np*ntot*sizeof(real_type)));
  real_type *Q_old = static_cast<real_type*>(malloc(np*ntot*sizeof(real_type)));
#else
  real_type *Q = new real_type[np*ntot];
  real_type *rhs = new real_type[np*ntot];
  real_type *Q_tmp = new real_type[np*ntot];
  real_type *Q_old = new real_type[np*ntot];
#endif

  cout << "BW CPU time is for 4 x iters for iters=" <<niters << endl;
  cout << "PHASEC" << endl;
  //mambo::begin();
  Timer wtimer;
  double t0=0,t1=0,t2=0;
#pragma omp parallel reduction(+:t0,t1,t2)
  {
    int first=ntot*omp_get_thread_num();
    int last=first+ntot;
    std::fill(Q+first,Q+last,0.0);
    std::fill(rhs+first,rhs+last,0.0);
    std::fill(Q_tmp+first,Q_tmp+last,0.0);
    std::fill(Q_old+first,Q_old+last,0.0);

    const double one_third=1.0/3.0;
    const double delta=0.01/2.0;
    vector<double> timing(4,0.0);

    //Q +=  one_third*rhs
    //Q_tmp=Q_old+delta*rhs
    Timer myclock;
    for(int i=0; i<niters; ++i)
    {
      myclock.restart();
      dns_ops::apply(Q+first,Q_tmp+first,Q_old+first,rhs+first,one_third,delta,ntot);
      timing[0]+=myclock.elapsed();

      myclock.restart();
      do_rk4(Q+first,Q_tmp+first,Q_old+first,rhs+first,&one_third,&delta,&ntot);
      timing[1]+=myclock.elapsed();

      myclock.restart();
      dns_ops::apply1(Q+first,rhs+first,one_third,ntot);
      dns_ops::apply2(Q_tmp+first,Q_old+first,rhs+first,delta,ntot);
      timing[2]+=myclock.elapsed();

      //myclock.restart();
      //dns_ops::apply(complex_ptr(Q),complex_ptr(Q_tmp),complex_ptr(Q_old),complex_ptr(rhs)
      //    ,one_third,delta,ntot/2);
      //timing[2]+=myclock.elapsed();
    }

    t0=timing[0];
    t1=timing[1];
    t2=timing[2];
  }

  double tot_time=wtimer.elapsed();
  //using 4 for RK4
  const double rk4=4.0;
  double nodefac= rk4/static_cast<double>(niters*omp_get_max_threads());
  cout.setf(ios::fixed, ios::floatfield);
  cout.precision(5);
  cout  << "BW OMP_NUM_THREADS = " << omp_get_max_threads() 
    <<  " Wall-clock      = " << tot_time << " sec " << endl;
  //cout << "fft-dim num-ffts ops loop1 loop2 loop1+complex " << endl;
  cout << "BW fft-dim num-ffts ops loop1 f77 " << endl;
  cout << "BW " << nz << " " << nc  << " " << nz*nc*n_var
    << " " << t0*nodefac << " " << t1*nodefac  << " " << t2*nodefac
    << endl;

  //mambo::end();
  cout << "PHASED" << endl;

#if defined(USE_MALLOC)
  free(Q);
  free(rhs);
  free(Q_tmp);
  free(Q_old);
#else
  delete [] Q;
  delete [] rhs;
  delete [] Q_tmp;
  delete [] Q_old;
#endif

  return 0;
}

